import pandas as pd
import os
from mongoengine import *
base_dir_00 = 'folder00'
base_dir_01 = 'folder01'
connect('ODMTest')


class Stock(Document):
    # 公司 証券代碼
    company = StringField(required=True, unique=True)
    # 简称
    abbreviation = StringField()
    # 上市交易所
    listed_exchange = StringField()
    # A股 上市码
    listing_code = StringField()
    # SSE产业分类
    sse = StringField()
    # 主要产品比重
    major_product_proportion = StringField()


def merge_mongodb(stk_id):
    stock = Stock()
    # print(stk_id)
    file_path_00 = base_dir_00 + '/' + stk_id+'.csv'
    file_path_01 = base_dir_01 + '/' + stk_id + '.csv'
    if os.path.exists(file_path_00):
        df_00 = pd.read_csv(file_path_00)
        coll_name = df_00['証券代碼'].values[0]
        print(coll_name)
        stock.switch_collection(str(coll_name))
        # for _, row in df_00.iteritems(): 遍历列
        # 遍历行
        for _, row in df_00.iterrows():
            # query first and the update
            stock.company = str(row['証券代碼'])
            stock.abbreviation = str(row['简称'])
            stock.listed_exchange = str(row['上市交易所'])
            stock.sse = str(row['SSE产业分类'])
            stock.major_product_proportion = str(row['主要产品比重'])
            stock.save()
            # print(stock)
            # exit()

    if os.path.exists(file_path_01):
        df_01 = pd.read_table(file_path_01, sep=';')
        # the same
        # print('df_01', df_01.shape)


def main():
    stk_ids_00 = os.listdir(base_dir_00)
    stk_ids_01 = os.listdir(base_dir_01)
    stk_ids = set(stk_ids_00+stk_ids_01)
    stk_ids = [stk_id.split('.')[0] for stk_id in stk_ids]
    # print(len(stk_ids))
    for i in range(5):
        merge_mongodb(stk_ids[i])


if __name__ == '__main__':
    main()
